"""

数据库执行公共SQL

"""
import re
import time

from testcase.utils.CommonSH import CommonSH


class CommonSQL(CommonSH):
    def __init__(self, node_name='dbuser'):
        """
        初始化
        :param node_name: 通过root用户还是数据库安装的用户执行脚本（见 conf/nodes.yml）
        """
        CommonSH.__init__(self, node_name)

    def create_func(self, func_name='', var='', execute_sql='', start='',
                    end='', step=''):
        """
        创建公用function，按照间隔值执行特定SQL
        return: func_name
        """
        function_cmd = f'''create or replace function {func_name}(num int)
            returns void as \$\$
            declare 
                {var} int;
            begin
                if num < {start} then 
                    for {var} in 0..num loop 
                        {execute_sql}
                    end loop;
                elsif num >= {start} and num < {start}+{step}*1 then
                    for {var} in {start}..num loop 
                        {execute_sql}
                    end loop;
                elsif num >= {start}+{step}*1 and num < {start}+{step}*2 then
                    for {var} in {start}+{step}*1..num loop 
                        {execute_sql}
                    end loop;
                elsif num >= {start}+{step}*2 and num < {end} then
                    for {var} in {start}+{step}*2..num loop 
                        {execute_sql}
                    end loop;
                end if;
            end;
            \$\$ language plpgsql;'''
        self.log.info(function_cmd)
        function_msg = self.execut_db_sql(function_cmd)
        self.log.info(function_msg)
        if function_msg.find(self.Constant.CREATE_FUNCTION_SUCCESS_MSG) > -1:
            return func_name
        else:
            self.log.info('---创建函数失败---')
            return False

    def create_vector_insert_func(self, func_name=''):
        """
        创建数据datavec插入函数，可自定义维度和插入记录数量、表名和列名及数据类型
        数据类型包括：vector、sparsevec、bit
        :param func_name: 函数名
        :return: func_name
        """
        create_fun_sql = f'''create or replace function {func_name}
            (vector_dim int,record_num int,table_name varchar,
            column_name varchar,date_type varchar) 
            returns varchar language 'plpgsql' as
            \$\$
            declare
                vector_value float[];
                i int;
                j int;
                v_sql varchar:='';
            begin
                for j in 1..record_num loop
                IF date_type = 'vector' THEN
                    for i in 1..vector_dim loop
                        vector_value[i] := round(random() * 100, 2);
                    end loop;
                    v_sql := format('insert into %I (%I) values (%L)',
                    table_name, column_name, vector_value::vector);
                ELSEIF date_type ='sparsevec' THEN
                    for i in 1..vector_dim loop
                        vector_value[i] := round(random() * 100, 2);
                    end loop;                
                    v_sql := format('insert into %I (%I) values (%L)',
                    table_name, column_name, vector_value::vector::sparsevec);
                ELSE
                    v_sql := format('
                    INSERT INTO %I (%I) SELECT (SELECT string_agg(
                    substr(''01'', ceil(random() * 2)::int, 1), \'\'\'\') 
                    FROM generate_series(1, %s))::bit(%s)', 
                    table_name, column_name, vector_dim, vector_dim);
                END IF;
                execute v_sql;
                end loop;
                return 'SUCCESS';
            end;
            \$\$;'''
        function_msg = self.execut_db_sql(create_fun_sql)
        self.log.info(function_msg)
        if function_msg.find(self.Constant.CREATE_FUNCTION_SUCCESS_MSG) > -1:
            return func_name
        else:
            self.log.info('---创建函数失败---')
            return False

    def vector_query_recall(self, query_sql, set_param='', is_detail=True):
        """
        向量查询召回率
        :param query_sql: 查询sql，查询字段第一个字段为id,可唯一标识一条记录
        :param set_param: 执行向量索引查询sql前设置的参数
        :param is_detail: 是否打印sql详细执行结果
        :return: 召回率
        """
        pattern = r'-+?\n(.*)\(\d+\srows?\)'
        self.log.info('----使用索引查询----')
        param = "set enable_indexscan=on;" \
                "set enable_seqscan=off;" \
                "set enable_bitmapscan=off;"
        sql_res = self.execut_db_sql(param + set_param + query_sql)
        if is_detail:
            self.log.info(sql_res)
        res = re.search(pattern, sql_res, re.M | re.S)
        rows = [[j.strip() for j in i.split('|')] for i in
                res.group(1).splitlines()]
        cols = list(map(list, zip(*rows)))
        retrieved_ids = cols[0]
        self.log.info(retrieved_ids)
        self.log.info('----未使用索引查询----')
        param = f"set enable_indexscan=off;" \
                f"set enable_seqscan=on;"
        sql_res = self.execut_db_sql(param + query_sql)
        res = re.search(pattern, sql_res, re.M | re.S)
        if is_detail:
            self.log.info(sql_res)
        rows = [[j.strip() for j in i.split('|')] for i in
                res.group(1).splitlines()]
        cols = list(map(list, zip(*rows)))
        true_ids = cols[0]
        self.log.info(true_ids)
        recall = len(set(retrieved_ids) & set(true_ids)) / len(true_ids)
        self.log.info('召回率为：' + str(recall) + '\n')
        return recall

    def get_query_duration(self, query_sql, set_param='', is_detail=True):
        """
        获取sql执行时长及执行结果
        :param query_sql: 执行SQL语句
        :param set_param: 执行sql前设置的参数等操作
        :param is_detail: 是否打印sql详细执行结果
        :return: 执行时长，执行结果
        """
        begin_time = time.time()
        sql_res = self.execut_db_sql(set_param + query_sql)
        if is_detail:
            self.log.info(sql_res)
        end_time = time.time()
        delta_time = end_time - begin_time
        return float(delta_time), sql_res
